# -*- coding: utf-8 -*-
# @Time    : 2024/6/19 14:10
# @Author  : sjh
# @Site    : 
# @File    : export_trt.py
# @Comment :
import torch
import logging
import onnx
from pathlib import Path
LOGGER = logging.getLogger('segment')  # define globally (used in train.py, val.py, detect.py, etc.)


def export_engine(im, file, half, dynamic=False, simplify=False, workspace=4, verbose=False):
    """
    Exports a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0.

    https://developer.nvidia.com/tensorrt
    """
    assert im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. `python export.py --device 0`"
    try:
        import tensorrt as trt
    except Exception:

        import tensorrt as trt


    onnx = file.with_suffix(".onnx")
    LOGGER.info(f"\n 111 starting export with TensorRT {trt.__version__}...")
    is_trt10 = int(trt.__version__.split(".")[0]) >= 10  # is TensorRT >= 10
    assert onnx.exists(), f"failed to export ONNX file: {onnx}"
    f = file.with_suffix(".engine")  # TensorRT engine file
    logger = trt.Logger(trt.Logger.INFO)
    if verbose:
        logger.min_severity = trt.Logger.Severity.VERBOSE

    builder = trt.Builder(logger)
    config = builder.create_builder_config()
    if is_trt10:
        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30)
    else:  # TensorRT versions 7, 8
        config.max_workspace_size = workspace * 1 << 30
    flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    network = builder.create_network(flag)
    parser = trt.OnnxParser(network, logger)
    if not parser.parse_from_file(str(onnx)):
        raise RuntimeError(f"failed to load ONNX file: {onnx}")

    inputs = [network.get_input(i) for i in range(network.num_inputs)]
    outputs = [network.get_output(i) for i in range(network.num_outputs)]
    for inp in inputs:
        LOGGER.info(f'111 input "{inp.name}" with shape{inp.shape} {inp.dtype}')
    for out in outputs:
        LOGGER.info(f'111 output "{out.name}" with shape{out.shape} {out.dtype}')

    if dynamic:
        if im.shape[0] <= 1:
            LOGGER.warning(f"111 WARNING ⚠️ --dynamic model requires maximum --batch-size argument")
        profile = builder.create_optimization_profile()
        for inp in inputs:
            profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)
        config.add_optimization_profile(profile)

    LOGGER.info(f"111 building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}")
    if builder.platform_has_fast_fp16 and half:
        config.set_flag(trt.BuilderFlag.FP16)

    build = builder.build_serialized_network if is_trt10 else builder.build_engine
    with build(network, config) as engine, open(f, "wb") as t:
        t.write(engine if is_trt10 else engine.serialize())

device = torch.device('cuda:0')
im = torch.zeros(1, 3, 640, 640).to(device)  # image size(1,3,320,192) BCHW iDetection
export_engine(im, file=Path("../weights/yolov10s.onnx"), half=True)
export_engine(im, file=Path("../weights/end2end.onnx"), half=False)